#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <unistd.h>
#include <uv.h>

#define UVTEE_BUF_SIZE 128

struct uvtee_context {
    uv_loop_t *loop;
    uv_pipe_t stdin_pipe;
    uv_pipe_t stdout_pipe;
    uv_pipe_t file_pipe;
    char *read_buffer;        // 预分配的读取缓冲区
    size_t buffer_size;       // 当前缓冲区大小
};

static void uvtee_context_init(struct uvtee_context *ctx, uv_loop_t *loop, int fd) {
    ctx->loop = loop;
    ctx->buffer_size = UVTEE_BUF_SIZE;
    ctx->read_buffer = malloc(ctx->buffer_size);
    
    uv_pipe_init(ctx->loop, &ctx->stdin_pipe, 0);
    uv_pipe_open(&ctx->stdin_pipe, STDIN_FILENO);
    uv_pipe_init(ctx->loop, &ctx->stdout_pipe, 0);
    uv_pipe_open(&ctx->stdout_pipe, STDOUT_FILENO);
    uv_pipe_init(ctx->loop, &ctx->file_pipe, 0);
    uv_pipe_open(&ctx->file_pipe, fd);
    ctx->stdin_pipe.data = ctx;
    ctx->stdout_pipe.data = ctx;
    ctx->file_pipe.data = ctx;
}

static void uvtee_context_cleanup(struct uvtee_context *ctx) {
    if (ctx->read_buffer) {
        free(ctx->read_buffer);
        ctx->read_buffer = NULL;
    }
    uv_close((uv_handle_t *)&ctx->stdin_pipe, NULL);
    uv_close((uv_handle_t *)&ctx->stdout_pipe, NULL);
    uv_close((uv_handle_t *)&ctx->file_pipe, NULL);
}

static void alloc_cb(uv_handle_t *handle, size_t suggested_size, uv_buf_t *buf) {
    struct uvtee_context *ctx = handle->data;
    
    // 如果建议的大小大于当前缓冲区，重新分配
    if (suggested_size > ctx->buffer_size) {
        char *new_buffer = realloc(ctx->read_buffer, suggested_size);
        if (new_buffer) {
            ctx->read_buffer = new_buffer;
            ctx->buffer_size = suggested_size;
        }
    }
    
    // 使用预分配的缓冲区
    *buf = uv_buf_init(ctx->read_buffer, ctx->buffer_size);
}

static void write_data(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf);
static void read_cb(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf) {
    struct uvtee_context *ctx = stream->data;
    if (nread < 0) {
        if (nread != UV_EOF) {
            fprintf(stderr, "Error reading from stdin: %s\n", uv_strerror(nread));
        } else {
            fprintf(stderr, "Read end of file\n");
        }
        uv_read_stop(stream);
    } else if (nread > 0) {
        write_data((uv_stream_t *)&ctx->stdout_pipe, nread, buf);
        write_data((uv_stream_t *)&ctx->file_pipe, nread, buf);
    } else {
        uv_read_stop(stream);
    }
    // 注意：不再释放buf->base，因为使用的是预分配的缓冲区
}

static void write_cb(uv_write_t *req, int status) {
    if (status < 0 || !req) {
        fprintf(stderr, "Error writing to stdout: %s\n", uv_strerror(status));
    }
    if (req->bufs && req->bufs[0].base) {
        free(req->bufs[0].base);
    }
    if (req->bufs) {
        free(req->bufs);
    }
    free(req);
}

static void write_data(uv_stream_t *stream, ssize_t nread, const uv_buf_t *buf) {
    uv_write_t *write_req = malloc(sizeof(uv_write_t));
    write_req->bufs = malloc(sizeof(uv_buf_t));
    write_req->bufs[0] = uv_buf_init(malloc(nread), nread);
    memcpy(write_req->bufs[0].base, buf->base, nread);
    uv_write(write_req, stream, write_req->bufs, 1, write_cb);
}

int main(int argc, char **argv) {
    if (argc < 2) {
        fprintf(stderr, "Usage: %s <file>\n", argv[0]);
        return 1;
    }

    uv_loop_t *loop = uv_default_loop();
    uv_fs_t open_req;
    int fd = uv_fs_open(loop, &open_req, argv[1], O_CREAT | O_TRUNC | O_WRONLY, 0644, NULL);
    if (fd < 0) {
        fprintf(stderr, "Error opening file: %s\n", uv_strerror(fd));
        return 1;
    }

    struct uvtee_context ctx;
    uvtee_context_init(&ctx, loop, fd);

    uv_read_start((uv_stream_t *)&ctx.stdin_pipe, alloc_cb, read_cb);

    int ret = uv_run(ctx.loop, UV_RUN_DEFAULT);

    uv_fs_close(loop, &open_req, fd, NULL);
    uvtee_context_cleanup(&ctx);
    uv_loop_close(loop);

    return ret;
}